import itertools
import copy
import pandas as pd
import statsmodels.api as sm
import pingouin as pg
from utils.analysis.common.mediation_common_class import RegressionDetails, RegressionModel, Model4Mediation, \
    Model6Result


def cal_f_value(cur_df: pd.DataFrame, x_cols: [], y_col: str):
    f_df = cur_df
    # 计算F值
    cs_name = '常量'  # 初始化常量
    f_df[cs_name] = 1
    # 提取自变量和因变量, 自变量，中介变量和控制变量合并为自变量
    # f_x = f_df[[cs_name, 'JG', "性别", "年龄", "学历", "企业角色"]]
    # f_y = f_df['YY']
    f_x = f_df[x_cols]
    f_y = f_df[y_col]
    # 添加截距项
    f_x = sm.add_constant(f_x)
    # 创建多元线性回归模型
    f_result = sm.OLS(f_y, f_x).fit()
    cur_res = dict()
    cur_res['f'] = f_result.fvalue
    cur_res['p'] = f_result.f_pvalue
    return cur_res


def analysis_model6(cur_df: pd.DataFrame, x, y, m1, m2, covar):
    # 返回结果：1. M1 <- X；2. M2 <- X + M1; 3. Y <- X; 4. Y <- X + M1 + M2
    mediation_res1 = pg.mediation_analysis(data=cur_df, x=x, m=[m1, m2], y=y, covar=covar, seed=0, n_boot=5000)
    # print(mediation_res1)
    mediation_res2 = pg.mediation_analysis(data=cur_df, x=x, m=[m1], y=y, covar=covar, seed=0, n_boot=5000)
    # print(mediation_res2)
    mediation_res3 = pg.mediation_analysis(data=cur_df, x=x, m=[m2], y=y, covar=covar, seed=0, n_boot=5000)
    # print(mediation_res3)

    base_x_arr = [x]
    if covar:
        base_x_arr.extend(covar)

    # m1 <- x
    x1 = copy.deepcopy(base_x_arr)
    m1_x = pg.linear_regression(cur_df[x1], cur_df[m1])
    m1_x_fp = cal_f_value(cur_df, x1, m1)

    # m2 <- x + m1
    x2 = copy.deepcopy(base_x_arr)

    x2.append(m1)
    m2_x_m1 = pg.linear_regression(cur_df[x2], cur_df[m2])
    m2_x_m1_fp = cal_f_value(cur_df, x2, m2)

    # y <- x
    x3 = copy.deepcopy(base_x_arr)
    y_x = pg.linear_regression(cur_df[x3], cur_df[y])
    y_x_fp = cal_f_value(cur_df, x3, y)

    # y <- x + m1 + m2
    x4 = copy.deepcopy(base_x_arr)
    x4.append(m1)
    x4.append(m2)
    y_x_m1_m2 = pg.linear_regression(cur_df[x4], cur_df[y])
    y_x_m1_m2_fp = cal_f_value(cur_df, x4, y)

    cur_res = dict()
    cur_res['effect_analysis_table'] = mediation_res1
    cur_res['regression_m1_x'] = m1_x
    cur_res['regression_m1_x_fp'] = m1_x_fp
    cur_res['regression_m2_x_m1'] = m2_x_m1
    cur_res['regression_m2_x_m1_fp'] = m2_x_m1_fp
    cur_res['regression_y_x'] = y_x
    cur_res['regression_y_x_fp'] = y_x_fp
    cur_res['regression_y_x_m1_m2'] = y_x_m1_m2
    cur_res['regression_y_x_m1_m2_fp'] = y_x_m1_m2_fp
    return cur_res


class MediationModel6:
    def __init__(self, df: pd.DataFrame, x_arr: [], y_arr: [], m_arr: [], covar: [], n_boot=5000):
        self.df = df
        self.x_arr = x_arr
        self.y_arr = y_arr
        self.m_arr = m_arr
        self.covar = covar
        self.n_boot = n_boot

    def item_analysis(self, x, y, m1, m2):
        cur_df = self.df
        # 返回结果：1. M1 <- X；2. M2 <- X + M1; 3. Y <- X; 4. Y <- X + M1 + M2
        mediation_res1 = pg.mediation_analysis(data=cur_df, x=x, m=[m1, m2], y=y, covar=self.covar, seed=0,
                                               n_boot=self.n_boot)

        base_x_arr = [x]
        if self.covar:
            base_x_arr.extend(self.covar)

        # m1 <- x
        x1 = copy.deepcopy(base_x_arr)
        m1_x = pg.linear_regression(cur_df[x1], cur_df[m1])
        m1_x_fp = cal_f_value(cur_df, x1, m1)

        # m2 <- x + m1
        x2 = copy.deepcopy(base_x_arr)
        x2.append(m1)
        m2_x_m1 = pg.linear_regression(cur_df[x2], cur_df[m2])
        m2_x_m1_fp = cal_f_value(cur_df, x2, m2)

        # y <- x
        x3 = copy.deepcopy(base_x_arr)
        y_x = pg.linear_regression(cur_df[x3], cur_df[y])
        y_x_fp = cal_f_value(cur_df, x3, y)

        # y <- x + m1 + m2
        x4 = copy.deepcopy(base_x_arr)
        x4.append(m1)
        x4.append(m2)
        y_x_m1_m2 = pg.linear_regression(cur_df[x4], cur_df[y])
        y_x_m1_m2_fp = cal_f_value(cur_df, x4, y)

        # 表
        mediation_table = [Model4Mediation(*row) for row in mediation_res1.to_numpy()]
        # 模型1
        model_m1x = RegressionModel(f=m1_x_fp.get('f'), p=m1_x_fp.get('p'),
                                    details=[RegressionDetails(*row) for row in m1_x.to_numpy()])
        # 模型2
        model_m2xm1 = RegressionModel(f=m2_x_m1_fp.get('f'), p=m2_x_m1_fp.get('p'),
                                      details=[RegressionDetails(*row) for row in m2_x_m1.to_numpy()])
        # 模型3
        model_yx = RegressionModel(f=y_x_fp.get('f'), p=y_x_fp.get('p'),
                                   details=[RegressionDetails(*row) for row in y_x.to_numpy()])
        # 模型4
        model_yxm1m2 = RegressionModel(f=y_x_m1_m2_fp.get('f'), p=y_x_m1_m2_fp.get('p'),
                                       details=[RegressionDetails(*row) for row in y_x_m1_m2.to_numpy()])
        # 总返回结果
        res_obj = Model6Result(mediation_table=mediation_table, m1_x=model_m1x, m2_x_m1=model_m2xm1, y_x=model_yx,
                               y_x_m1_m2=model_yxm1m2, x=x, y=y, m1=m1, m2=m2, covar=self.covar)
        return res_obj

    def analysis(self):
        res = []
        combinations = list(itertools.combinations(self.m_arr, 2))
        for x in self.x_arr:
            for y in self.y_arr:
                for m in combinations:
                    item_res = self.item_analysis(x, y, m[0], m[1])
                    res.append(item_res)
        return res


if __name__ == '__main__':
    pd.set_option('expand_frame_repr', False)
    # 逐步回归做
    # 链式中介
    # 如果有多个中介变量，每一次回归需要取两个中介，放到自变量，控制变量也放到自变量
    # 返回结果：1. M1 <- X；2. M2 <- X + M1; 3. Y <- X; 4. Y <- X + M1 + M2

    # 1. 链式中介

    df = pd.read_excel('./test_datas.xlsx')
    df = df[["性别", "年龄", "学历", "企业角色", 'JG', 'YY', 'JX', 'XW']]
    # 1.此文件的结果不包含整体的F值和p值
    # 2.若有多个x,y,m，则需要遍历进行获取结果
    x_arr = ['JG']
    y_arr = ['YY']
    m_arr = ['XW', 'JX']
    covar = ["性别", "年龄", "学历", "企业角色"]
    # covar = []
    res = []
    combinations = list(itertools.combinations(m_arr, 2))
    for x in x_arr:
        for y in y_arr:
            for m in combinations:
                item_res = analysis_model6(df, x, y, m[0], m[1], covar)
                res.append(item_res)
                for cc in item_res:
                    print(cc)
                    print(item_res[cc])
    print('===================================================================================================')
